import json
import asyncio
import time
import struct

import uvicorn
from fastapi import FastAPI, HTTPException, status
from fastapi.responses import HTMLResponse, StreamingResponse
import cv2
import numpy as np

from logger import log

app = FastAPI()

# 存储视频帧
"""
{
    "vs1": {
        "client_socket": client_socket
        "streams": {
            id:{
                "name": "name",
                "buf": buf,
                "count": 0,
            }
        },
        "last_heartbeat": time.time()
    },
}
"""
edges = {}
edges_lock = asyncio.Lock()
HEARTBEAT_INTERVAL = 30
HEARTBEAT_TIMEOUT = 60


async def get_frame(edge_code, id):
    while True:
        async with edges_lock:
            if edge_code not in edges:
                break
            frame = edges[edge_code]['streams'][id].get('frame')
        if frame is not None:
            ret, buffer = cv2.imencode('.jpg', frame)
            frame_bytes = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + frame_bytes + b'\r\n')
        else:
            yield b''
        await asyncio.sleep(0.2)


@app.get("/", response_class=HTMLResponse)
async def hello():
    async with edges_lock:
        links = []
        for edge_key, edge in edges.items():
            streams = edge['streams']
            for key, stream in streams.items():
                links.append(f'<li><a href="/video/{edge_key}/{key}">{stream["name"]}:{stream["count"]}</a></li>')
        links_html = '<ul>' + ''.join(links) + '</ul>'
    return f'<html><body><h1>Streams</h1>{links_html}</body></html>'


@app.get("/{edge_code}/{id_str}")
async def video_feed(edge_code: str, id_str: str):
    log.info(f"video_feed:{edge_code},{id_str}")
    id = int(id_str)
    async with edges_lock:
        if edge_code not in edges:
            log.info(f"没有找到客户端:{edge_code}")
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Edge code '{edge_code}' not found")
        client = edges[edge_code]
        if id not in client['streams']:
            raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Edge code id '{id}' not found")
        client['streams'][id]['count'] += 1
        if client['streams'][id]['count'] == 1:
            log.info(f"发送连接请求:{id},stream['count']:{client['streams'][id]['count']}")
            await send_command(client['client_writer'], {"cmd": "start", "id": id})

    async def generator():
        try:
            async for frame in get_frame(edge_code, id):
                if frame:
                    yield frame
        except asyncio.CancelledError:
            log.info(f"GeneratorExit triggered for stream {id}")
        finally:
            async with edges_lock:
                client['streams'][id]['count'] -= 1
                log.info(f"客户端退出:{id},stream['count']:{client['streams'][id]['count']}")
                if client['streams'][id]['count'] == 0:
                    log.info(f"发送关闭连接请求:{id}")
                    await send_command(client['client_writer'], {"cmd": "stop", "id": id})

    return StreamingResponse(generator(), media_type='multipart/x-mixed-replace; boundary=frame')


async def send_command(writer, command):
    try:
        data = json.dumps(command).encode()
        # Assuming the client expects the command size first
        writer.write(struct.pack('!I', len(data)) + data)
        await writer.drain()
    except Exception as e:
        log.exception(e)


async def handle_client_connection(reader, writer):
    peername = writer.get_extra_info('peername')
    log.info(f"handle_client_connection: {peername}")
    payload_size = struct.calcsize("QQ")

    async def heartbeat_check(edge_code):
        while True:
            await asyncio.sleep(HEARTBEAT_INTERVAL)
            async with edges_lock:
                if edge_code not in edges:
                    break
                last = edges[edge_code]['last_heartbeat']
                if time.time() - last > HEARTBEAT_TIMEOUT:
                    log.warning(f"Heartbeat timeout for {edge_code}, disconnecting...")
                    log.info(f"del edges1:{edge_code}")
                    if not writer.is_closing():
                        writer.close()
                        try:
                            await writer.wait_closed()
                        except Exception as e_wc_main:
                            log.error(
                                f"Error during wait_closed for main writer of {edge_code or 'unknown_edge'}: {e_wc_main}")
                    del edges[edge_code]
                    log.info(f"del edges2:{edge_code}")
                    break

    try:
        # 读取4字节数据长度
        length_data = await reader.readexactly(4)
        data_length = struct.unpack('!I', length_data)[0]

        # 读取完整JSON数据
        data = await reader.readexactly(data_length)
        obj = json.loads(data.decode())
        edge_code = obj['edge_code']
        streams = {stream['id']: {"name": stream['name'], "frame": None, "count": 0} for stream in obj['streams']}

        async with edges_lock:
            edges[edge_code] = {
                "client_writer": writer,
                "streams": streams,
                "last_heartbeat": time.time()
            }

        # 启动心跳检测任务
        asyncio.create_task(heartbeat_check(edge_code))

        while True:
            try:
                delimiter = await reader.readexactly(4)
                if delimiter != b'@#@#':
                    log.warning("Invalid delimiter received")
                    continue

                header = await reader.readexactly(payload_size)
                id, frame_size = struct.unpack("QQ", header)

                if id == 0 and frame_size == 0:
                    # 心跳包
                    async with edges_lock:
                        if edge_code in edges:
                            edges[edge_code]['last_heartbeat'] = time.time()
                    await send_command(writer, {"cmd": "heartbeat"})
                    continue

                # 读取帧数据
                frame_data = await reader.readexactly(frame_size)
                frame = np.frombuffer(frame_data, dtype=np.uint8)
                frame = cv2.imdecode(frame, cv2.IMREAD_COLOR)

                async with edges_lock:
                    if edge_code in edges and id in edges[edge_code]['streams']:
                        edges[edge_code]['streams'][id]['frame'] = frame
            except asyncio.IncompleteReadError:
                log.info(f"Connection closed by client: {peername}")
                break
            except Exception as e:
                log.exception(e)
                break
    except Exception as e:
        log.exception(e)

    log.info(f"Connection closed for {peername}")
    if not writer.is_closing():
        writer.close()
        try:
            await writer.wait_closed()
        except Exception as e_wc_main:
            log.error(f"Error during wait_closed for main writer of {edge_code or 'unknown_edge'}: {e_wc_main}")
    log.info(f"del edges1:{edge_code}")
    async with edges_lock:
        if edge_code in edges:
            del edges[edge_code]
            log.info(f"del edges2:{edge_code}")


async def start_tcp_server():
    server = await asyncio.start_server(handle_client_connection, '0.0.0.0', 9999)
    addr = server.sockets[0].getsockname()
    log.info(f'TCP server listening on {addr}')

    async with server:
        await server.serve_forever()


@app.on_event("startup")
async def startup_event():
    asyncio.create_task(start_tcp_server())


def start():
    # 使用 uvicorn 运行 FastAPI 应用
    # host="0.0.0.0" 表示监听所有网络接口，允许局域网内其他设备访问
    uvicorn.run(app, host="0.0.0.0", port=8888)


if __name__ == "__main__":
    start()

# 运行FastAPI应用时，建议使用uvicorn并启用多工作线程以处理高并发
# 可以使用以下命令运行应用：
# pip install fastapi uvicorn
# uvicorn your_module_name:app --host 0.0.0.0 --port 8888

# 定时重启cron crontab -e
# 0 3 * * * /root/restart-video.sh
# #!/bin/bash
#
# # 定义你要写入的文件
# OUTPUT_FILE="execution_times.log"
#
# # 获取当前时间并格式化
# current_time=$(date +"%Y-%m-%d %H:%M:%S")
#
# # 将当前时间写入文件
# echo "Executed at: $current_time" >> $OUTPUT_FILE
#
# # 这里可以加上你需要执行的其他命令
# systemctl restart video
